# coding = utf-8

'''
关于肿瘤分割的后处理
'''

import numpy as np
from scipy import ndimage
from skimage import measure

def select_tumor(output, label, prob):
    predict_liver = np.zeros(output.shape)
    predict_tumor = np.zeros(output.shape)
    predict_liver[output>=1] = 1
    predict_tumor[output == 2] = 1

    #select the biggest liver
    #predict_liver = ndimage.binary_dilation(predict_liver, iterations=1).astype(predict_liver.dtype)
    [liver_labels, num] = measure.label(predict_liver, return_num=True)
    region = measure.regionprops(liver_labels)
    box = []
    for i in range(num):
        box.append(region[i].area)
    label_num = box.index(max(box)) + 1
    liver_labels[liver_labels != label_num] = 0
    liver_labels[liver_labels == label_num] = 1
    #predict_liver = ndimage.binary_fill_holes(liver_labels).astype(int)

    #select tumor inside the liver
    label_copy = np.zeros(label.shape)
    label_copy[label >= 1] = 1
    predict_tumor = predict_tumor * label_copy


    #select tumor biggest
    real_tumor = np.zeros(predict_tumor.shape)

    #select tumor by prob each pixel

    '''
    prob = prob * label_copy
    real_tumor[prob>=0.6] = 1
    '''


    #predict_tumor = ndimage.binary_fill_holes(predict_tumor).astype(int)
    #select tumor by max prob

    [tumor_labels, num] = measure.label(predict_tumor, return_num=True)
    region = measure.regionprops(tumor_labels)
    box = []
    for i in range(num):
        tool_array = np.zeros(tumor_labels.shape)
        tool_array[tumor_labels == i+1] = 1
        tool_array = tool_array * prob
        box.append(np.max(tool_array))
        if np.max(np.array(tool_array)) >= 0.8:
            real_tumor[tumor_labels == i+1] =1
    print(sorted(box))


    '''
    label_copy = np.zeros(label.shape)
    label_copy[label == 2] = 1
    [_, num2] = measure.label(label_copy, return_num=True)
    print(num, num2, np.max(prob), np.min(prob))
    '''

    predict_liver[real_tumor == 1] = 2
    return predict_liver